import os
import random
import json
import paddle
import sys
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


# 定义公共变量
name_dict = {"apple": 0, "banana": 1, "grape": 2, "orange": 3, "pear": 4}
data_root_path = "fruits/" # 数据集目录
test_file_path = data_root_path + "test.txt" # 测试集文件路径
train_file_path = data_root_path + "train.txt" # 测试集文件
name_data_list = {} # 记录每个类别图片 key:名称  value:路径列表

def save_train_test_file(path, name): # 将图片添加到字典
    if name not in name_data_list: # 该类别水果不在字典中
        img_list = []
        img_list.append(path) # 路径存入列表
        name_data_list[name] = img_list # 列表存入字典
    else:
        name_data_list[name].append(path) # 直接添加到列表

# 遍历每个子目录，将图片路径存入字典
dirs = os.listdir(data_root_path) # 列出数据集下的子目录
for d in dirs:
    full_path = data_root_path + d # 子目录完整路径
    if os.path.isdir(full_path): # 如果是目录
        imgs = os.listdir(full_path) # 列出子目录下的图片
        for img in imgs:
            img_full_path = full_path + "/" + img # 图片路径
            save_train_test_file(img_full_path, d) # 添加到字典
    else: # 文件
        pass

# 划分训练集、测试集
with open(test_file_path, "w") as f:
    pass
with open(train_file_path, "w") as f:
    pass

# 遍历字典
for name, img_list in name_data_list.items():
    i = 0
    num = len(img_list) # 取出样本数量
    print("%s: %d张图像" % (name, num))

    for img in img_list:
        # 拼接一行
        line = "%s\t%d\n" % (img, name_dict[name])
        if i % 10 == 0: # 写入测试集
            with open(test_file_path, "a") as f:
                f.write(line) # 存入文件
        else: # 写入训练集
            with open(train_file_path, "a") as f:
                f.write(line) # 存入文件
        i += 1
print("数据预处理完成.")
